"""
高级个人助理 - 完整记忆管理示例

这个版本包含了更完整的记忆管理功能，包括：
1. 多种类型的偏好管理
2. 智能上下文理解
3. 任务追踪和提醒
4. 更好的错误处理
"""

import json
import datetime
from typing import Optional, Dict, Any, List
from google.adk.agents import Agent
from google.adk.tools import ToolContext

import os
from google.adk.models.lite_llm import LiteLlm

# API密钥配置
DEFAULT_DASHSCOPE_API_KEY = "sk-f227634bb56*************************232"  # 请替换为您的实际密钥
DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY", DEFAULT_DASHSCOPE_API_KEY)

# 常量定义
APP_NAME = "sequential_document_processor"
USER_ID = "dev_user_01"
GEMINI_MODEL = LiteLlm(
    model="openai/qwen-turbo",  # 使用通义千问Turbo模型
    api_key=DASHSCOPE_API_KEY,
    api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
)


def save_user_preference(tool_context: ToolContext, preference_name: str, preference_value: str, category: str = "general") -> str:
    """
    保存用户偏好，支持分类管理。
    
    Args:
        tool_context: 工具上下文对象
        preference_name: 偏好设置的名称
        preference_value: 偏好设置的值
        category: 偏好分类（如 'personal', 'work', 'food', 'travel'）
        
    Returns:
        确认消息
    """
    try:
        # 保存具体偏好
        state_key = f"user:pref_{preference_name}"
        tool_context.state[state_key] = {
            "value": preference_value,
            "category": category,
            "updated_at": datetime.datetime.now().isoformat()
        }
        
        # 更新偏好列表
        preferences_key = "user:saved_preferences"
        saved_preferences = tool_context.state.get(preferences_key, {})
        if category not in saved_preferences:
            saved_preferences[category] = []
        if preference_name not in saved_preferences[category]:
            saved_preferences[category].append(preference_name)
        tool_context.state[preferences_key] = saved_preferences
        
        return f"好的，我已经记住您在{category}类别下的{preference_name}是{preference_value}。"
    
    except Exception as e:
        return f"抱歉，保存偏好时遇到了问题：{str(e)}"


def recall_user_info(tool_context: ToolContext, info_key: str) -> str:
    """
    智能回忆用户信息，支持模糊搜索。
    """
    try:
        # 直接查找
        direct_key = f"user:pref_{info_key}"
        direct_value = tool_context.state.get(direct_key)
        
        if direct_value:
            if isinstance(direct_value, dict):
                return f"我记得您的{info_key}是：{direct_value['value']} (分类：{direct_value.get('category', '未分类')})"
            else:
                return f"我记得您的{info_key}是：{direct_value}"
        
        # 模糊搜索
        all_prefs = {}
        for key, value in tool_context.state.items():
            if key.startswith("user:pref_"):
                pref_name = key[10:]  # 移除 "user:pref_" 前缀
                if info_key.lower() in pref_name.lower():
                    all_prefs[pref_name] = value
        
        if all_prefs:
            results = []
            for pref_name, pref_data in all_prefs.items():
                if isinstance(pref_data, dict):
                    results.append(f"- {pref_name}: {pref_data['value']}")
                else:
                    results.append(f"- {pref_name}: {pref_data}")
            return f"我找到了相关信息：\n" + "\n".join(results)
        
        # 尝试搜索长期记忆
        try:
            search_results = tool_context.search_memory(query=f"用户 {info_key}")
            if search_results and hasattr(search_results, 'results') and search_results.results:
                return f"从过去的对话中，我找到了关于{info_key}的信息：{search_results.results[0] if search_results.results else '未找到相关信息'}"
        except Exception:
            pass
        
        return f"抱歉，我还没有记录关于{info_key}的信息。您可以告诉我，我会记住的。"
    
    except Exception as e:
        return f"查询信息时遇到了问题：{str(e)}"


def list_preferences_by_category(tool_context: ToolContext, category: Optional[str] = None) -> str:
    """
    按分类列出用户偏好。
    """
    try:
        preferences_key = "user:saved_preferences"
        saved_preferences = tool_context.state.get(preferences_key, {})
        
        if not saved_preferences:
            return "我还没有保存任何关于您的偏好信息。"
        
        if category:
            if category not in saved_preferences:
                return f"没有找到{category}分类的偏好设置。"
            
            prefs_list = []
            for pref_name in saved_preferences[category]:
                state_key = f"user:pref_{pref_name}"
                pref_data = tool_context.state.get(state_key, "未设置")
                if isinstance(pref_data, dict):
                    prefs_list.append(f"- {pref_name}: {pref_data['value']}")
                else:
                    prefs_list.append(f"- {pref_name}: {pref_data}")
            
            return f"{category}分类的偏好设置：\n" + "\n".join(prefs_list)
        
        else:
            all_prefs = []
            for cat, prefs in saved_preferences.items():
                all_prefs.append(f"\n【{cat}】")
                for pref_name in prefs:
                    state_key = f"user:pref_{pref_name}"
                    pref_data = tool_context.state.get(state_key, "未设置")
                    if isinstance(pref_data, dict):
                        all_prefs.append(f"  - {pref_name}: {pref_data['value']}")
                    else:
                        all_prefs.append(f"  - {pref_name}: {pref_data}")
            
            return "您的所有偏好设置：" + "\n".join(all_prefs)
    
    except Exception as e:
        return f"获取偏好列表时遇到了问题：{str(e)}"


def add_task_reminder(tool_context: ToolContext, task_description: str, due_date: Optional[str] = None, priority: str = "medium") -> str:
    """
    添加任务提醒。
    """
    try:
        tasks_key = "user:tasks"
        tasks = tool_context.state.get(tasks_key, [])
        
        task = {
            "id": len(tasks) + 1,
            "description": task_description,
            "due_date": due_date,
            "priority": priority,
            "status": "pending",
            "created_at": datetime.datetime.now().isoformat()
        }
        
        tasks.append(task)
        tool_context.state[tasks_key] = tasks
        
        due_info = f"，截止日期：{due_date}" if due_date else ""
        return f"好的，我已经为您添加了任务：{task_description} (优先级：{priority}{due_info})"
    
    except Exception as e:
        return f"添加任务时遇到了问题：{str(e)}"


def list_tasks(tool_context: ToolContext, status: str = "all") -> str:
    """
    列出任务列表。
    """
    try:
        tasks_key = "user:tasks"
        tasks = tool_context.state.get(tasks_key, [])
        
        if not tasks:
            return "您目前没有任何任务。"
        
        if status != "all":
            tasks = [task for task in tasks if task.get("status") == status]
        
        if not tasks:
            return f"没有找到状态为'{status}'的任务。"
        
        task_list = []
        for task in tasks:
            status_icon = "✓" if task.get("status") == "completed" else "○"
            due_info = f" (截止：{task.get('due_date')})" if task.get('due_date') else ""
            priority_icon = {"high": "🔴", "medium": "🟡", "low": "🟢"}.get(task.get('priority', 'medium'), "🟡")
            
            task_list.append(f"{status_icon} {priority_icon} {task['description']}{due_info}")
        
        return f"您的任务列表：\n" + "\n".join(task_list)
    
    except Exception as e:
        return f"获取任务列表时遇到了问题：{str(e)}"


def update_task_status(tool_context: ToolContext, task_id: int, new_status: str) -> str:
    """
    更新任务状态。
    """
    try:
        tasks_key = "user:tasks"
        tasks = tool_context.state.get(tasks_key, [])
        
        if not tasks:
            return "没有找到任何任务。"
        
        for task in tasks:
            if task.get("id") == task_id:
                old_status = task.get("status", "unknown")
                task["status"] = new_status
                task["updated_at"] = datetime.datetime.now().isoformat()
                tool_context.state[tasks_key] = tasks
                return f"任务'{task['description']}'的状态已从'{old_status}'更新为'{new_status}'。"
        
        return f"没有找到ID为{task_id}的任务。"
    
    except Exception as e:
        return f"更新任务状态时遇到了问题：{str(e)}"


def get_personalized_greeting(tool_context: ToolContext) -> str:
    """
    根据用户信息生成个性化问候。
    """
    try:
        # 获取用户姓名
        name_key = "user:pref_name"
        user_name = tool_context.state.get(name_key)
        if isinstance(user_name, dict):
            user_name = user_name.get('value')
        
        # 获取时间偏好
        timezone_key = "user:pref_timezone"
        timezone = tool_context.state.get(timezone_key)
        if isinstance(timezone, dict):
            timezone = timezone.get('value')
        
        # 构建问候语
        greeting_parts = []
        
        current_hour = datetime.datetime.now().hour
        if current_hour < 12:
            time_greeting = "早上好"
        elif current_hour < 18:
            time_greeting = "下午好"
        else:
            time_greeting = "晚上好"
        
        if user_name:
            greeting_parts.append(f"{time_greeting}，{user_name}！")
        else:
            greeting_parts.append(f"{time_greeting}！")
        
        # 检查待办任务
        tasks_key = "user:tasks"
        tasks = tool_context.state.get(tasks_key, [])
        pending_tasks = [task for task in tasks if task.get("status") == "pending"]
        
        if pending_tasks:
            greeting_parts.append(f"您还有{len(pending_tasks)}个待完成的任务。")
        
        return " ".join(greeting_parts) + " 我能为您做些什么？"
    
    except Exception as e:
        return f"生成个性化问候时遇到了问题：{str(e)}"


# 创建高级个人助理代理
root_agent = Agent(
    name="advanced_personal_assistant",
    model=GEMINI_MODEL,
    description="一个具有完整记忆管理、任务追踪和个性化服务的高级个人助理",
    instruction="""你是一个高级智能个人助理，具有强大的记忆管理和任务管理能力。

**核心能力：**

1. **分类偏好管理**：
   - 使用 'save_user_preference' 保存不同类别的用户偏好
   - 支持的分类：personal（个人）、work（工作）、food（饮食）、travel（旅行）、health（健康）等
   - 使用 'list_preferences_by_category' 查看特定分类的偏好

2. **智能信息检索**：
   - 使用 'recall_user_info' 进行智能搜索，支持模糊匹配
   - 自动从用户状态和长期记忆中检索相关信息

3. **任务管理**：
   - 使用 'add_task_reminder' 添加任务提醒
   - 使用 'list_tasks' 查看任务列表
   - 使用 'update_task_status' 更新任务状态

4. **个性化交互**：
   - 使用 'get_personalized_greeting' 生成个性化问候
   - 基于用户历史和偏好定制回复

**交互原则：**
- 主动询问并记录重要的用户信息
- 在提供建议时考虑用户的偏好和历史
- 定期提醒用户的待办任务
- 保持友好、专业和高效的服务态度
- 在不确定时主动澄清用户需求

**示例使用场景：**
- "记住我喜欢中餐，分类为饮食偏好"
- "提醒我明天下午3点开会"
- "我的工作相关偏好有哪些？"
- "将任务1标记为完成"

记住：你的目标是成为用户最可靠、最了解他们需求的智能助理。""",
    
    tools=[
        save_user_preference,
        recall_user_info,
        list_preferences_by_category,
        add_task_reminder,
        list_tasks,
        update_task_status,
        get_personalized_greeting
    ]
) 